import os
os.system("module load apps/anaconda3/latest")
import numpy as np
import wrf
import gsw
import xarray as xr
import netCDF4 as nc
import matplotlib.pyplot as plt
from scipy.interpolate import CubicSpline,interp1d
from multiprocessing import Pool

wrfoutpath="./wrfout_d01_2018-09-13_00_00_00"
dir_coawst="./"


def parallel(inputs):
    t,x,y = inputs
    print(f"{t} start",flush = True)
    nz,ny,nx = x.shape
    stdT_0m   = np.zeros((ny,nx))
    stdT_100m = np.zeros( (ny,nx) )
    for j in range(ny):
        for i in range(nx):
            cs = interp1d(x[:,j,i], y[:,j,i],fill_value="extrapolate")
            stdT_0m[j,i]   = cs(0)
            stdT_100m[j,i] = cs(-100)
    return t,stdT_0m,stdT_100m


#类封装后处理方法，代码易读性好，复用性强
class wrfout:
    def __init__(self,wrfoutpath):
        ###   basic   ###
        ncf=nc.Dataset(wrfoutpath)
        self.ncf=ncf
        self.lon=wrf.getvar(ncf,"XLONG")
        self.lat=wrf.getvar(ncf,"XLAT")
        self.p = wrf.getvar(ncf,"pressure")
        Times=wrf.getvar(ncf,"Times",timeidx=wrf.ALL_TIMES)
        self.Times=[str(tt)[:13] for tt in Times.values]
        
        # 读取全部时刻or单个时刻 ; 全部时刻业务用，单个时刻开发用
        t=wrf.ALL_TIMES
        # t=0
        
        ###   specific  ###
        # 1等压面位势高度
        self.gph =wrf.getvar(ncf,'geopt',timeidx=t)/9.81 
        # 2等压面温度
        self.T   =wrf.getvar(ncf,'tk',timeidx=t)
        # 3等压面相对湿度
        self.rh  =wrf.getvar(ncf,"rh",timeidx=t)
        # 4等压面U风
        self.u   =wrf.getvar(ncf,'ua',timeidx=t)
        # 5等压面V风
        self.v   =wrf.getvar(ncf,'va',timeidx=t)
        # 6地面气压
        self.psfc   =wrf.getvar(ncf,'PSFC',timeidx=t)
        # 7海平面气压
        self.slp    =wrf.getvar(ncf,'slp',timeidx=t)
        # 8:10米高度U风
        self.u10    =wrf.getvar(ncf,'U10',timeidx=t)
        # 9:10米高度V风
        self.v10    =wrf.getvar(ncf,'V10',timeidx=t)
        # 10:2米高度温度
        self.T2     =wrf.getvar(ncf,'T2',timeidx=t)
        # 11地面温度
        self.tsk    =wrf.getvar(ncf,'TSK',timeidx=t)
        # 12降水
        self.rain   =wrf.getvar(ncf,"RAINC",timeidx=t)+wrf.getvar(ncf,"RAINNC",timeidx=t)+\
                    wrf.getvar(ncf,"RAINSH",timeidx=t)
        # 13海雾
        self.haiwu  =self.cal_haiwu(t)
        # 16台风  25.2.16 暂时用最低海平面气压表示
        self.tf = np.min(np.min(self.slp,axis=1),axis=1) 
        # 17高空颠簸指数---用ELLROD指数
        self.gkdbzs = xr.DataArray( self.cal_gkdbzs(),dims=self.T.dims,coords=self.T.coords  )
        # 18积冰指数
        self.jbzs = xr.DataArray( self.cal_jbzs() ,dims=self.T.dims,coords=self.T.coords  )
        # 20海表温度
        self.sst = wrf.getvar(ncf,'SST',timeidx=t)
        # 31大气波导 - 用修正折射率暂时表示
        self.dmdz = xr.DataArray( self.cal_dmdz(t) ,dims=self.T.dims,coords=self.T.coords  )
        # self.bdtype = xr.DataArray(self.cal_dqbdtype(self.dmdz, self.gph) ,dims=self.tsk.dims,coords=self.tsk.coords  )
        


    # 计算大气波导，用修正折射率M暂代替大气波导    
    def cal_dmdz(self,t):
        # 计算修正折射率
        q=wrf.getvar(self.ncf,"QVAPOR",timeidx=t) # 水汽混合比 (kg/kg)
        T = self.T  # 温度 (K)
        P = self.p  # 气压 (hPa) 
        h = self.gph  # 高度 (m)
        epsilon = 0.622        # 常数
        R = 6371 * 1e3  # 地球半径 (m)
        e = (q * P) / (epsilon + (1 - epsilon) * q)
        N = (776 / T) * (P + (4810 * e) / T)        
        M = N + (h / R) * 1e6         # 计算修正大气折射率 M
        dmdz = np.gradient(M,axis=1)/np.gradient(h,axis=1)
        return dmdz


    def cal_dqbdtype(self,dMdz, z_coords, threshold=-0.1):
    
        # 初始化结果数组
        duct_types = np.full((dMdz.shape[0], dMdz.shape[2], dMdz.shape[3]), "None", dtype=object)
    
        # 遍历每个时空点
        for t in range(dMdz.shape[0]):
            print("dqbd type",t,dMdz.shape[0])
            for y in range(dMdz.shape[2]):
                for x in range(dMdz.shape[3]):
                    dMdz_slice = dMdz[t, :, y, x]
                    z_coords = self.gph[t,:,y,x]
                    # 检测波导层
                    duct_layers = []
                    in_duct = False
                    base_idx = None
    
                    for z_idx in range(len(z_coords)):
                        if dMdz_slice[z_idx] < threshold:
                            if not in_duct:
                                in_duct = True
                                base_idx = z_idx  # 波导底层索引
                        else:
                            if in_duct:
                                in_duct = False
                                top_idx = z_idx - 1  # 波导顶层索引
                                duct_layers.append((base_idx, top_idx))
    
                    # 处理最后一层
                    if in_duct:
                        top_idx = len(z_coords) - 1
                        duct_layers.append((base_idx, top_idx))
    
                    # 判断波导类型
                    if len(duct_layers) == 0:
                        duct_type = "None"
                    else:
                        # 选择最显著的波导层（例如厚度最大）
                        max_thickness = 0
                        selected_layer = None
                        for layer in duct_layers:
                            thickness = z_coords[layer[1]] - z_coords[layer[0]]
                            if thickness > max_thickness:
                                max_thickness = thickness
                                selected_layer = layer
    
                        # 判断是否为表面波导
                        if selected_layer[0] == 0:
                            duct_type = "Surface Duct"
                        else:
                            duct_type = "Elevated Duct"
    
                    # 保存结果
                    duct_types[t, y, x] = duct_type
    
        return duct_types

    
    #求海雾的算法；海雾定义是能见度低于1km，haiwu函数实际为计算能见度
    #具体使用了NOAA FSL算法（简单）。详情见文献：基于 WＲF 模式的青岛近海能见度算法比较研究
    def cal_haiwu(self,t):
        Td = wrf.getvar(self.ncf,'td',timeidx=t).values+273.15
        return 6000*1.609344*(self.T-Td)/self.rh**1.75
        
    
    
    
    # 计算高空颠簸指数--ELLROD指数  论文：中南地区航空器空中颠簸指数适用分析
    def cal_gkdbzs(self):
        '''
        EI < 8：无颠簸或轻度颠簸。
        8 ≤ EI < 12：中等强度颠簸。
        EI ≥ 12：强颠簸
        '''
        if self.u.ndim==3:
            axis_z=0 ; axis_y=1 ;axis_x=2
        elif self.u.ndim==4:
            axis_z=1 ; axis_y=2 ; axis_x=3
        # vertical sheer
        dz = np.gradient(self.gph , axis=axis_z)  # 高度梯度
        dx = self.ncf.getncattr("DX")
        dy = self.ncf.getncattr("DY")
        du_z = np.gradient(self.u,axis=axis_z)  
        dv_z = np.gradient(self.v,axis=axis_z)  
        vertical_shear = np.sqrt( (du_z/dz)**2 + (dv_z/dz)**2 )
        # 计算水平变形项
        du_x = np.gradient(self.u,axis=axis_x)  # U分量水平梯度
        dv_y = np.gradient(self.v,axis=axis_y)  # V分量水平梯度
        dv_x = np.gradient(self.v,axis=axis_x)  # V分量水平梯度
        du_y = np.gradient(self.u,axis=axis_y)  # U分量水平梯度
        horizontal_deformation = np.sqrt( (du_x/dx + dv_y/dy)**2 + (dv_x/dx + du_y/dy)**2 )
        # 计算ELLROD2指数（忽略了散度项）
        EL = vertical_shear * horizontal_deformation
        return EL
        
    
    #积冰指数 -- IC指数 论文：飞机积冰预报方法及应用场景研究
    def cal_jbzs(self):
        '''
        <0 无
        0-50 轻度积冰
        50-80 中度
        80+  重度
        '''
        TC = self.T-273.15
        IC = ((self.rh-50)*2) * (TC* (TC+14)/(-49))
        IC = np.where(IC>100,100,IC)
        return IC
        
    
        
        
    # 快速看一下数值，检查正确性;输入的data可以是气象的2，3，4维数据
    def see(self,data):
        if data.ndim == 2:
            plt.contourf(data)
            plt.colorbar()
            plt.show()
        elif data.ndim == 3:
            num_slices = data.shape[0]  # 第一个维度的大小
            for i in range(num_slices):
                plt.figure()
                plt.title(np.mean(self.p[i,:,:]))
                plt.imshow(data[i, :, :], cmap='viridis')  # 绘制第i个切片
                plt.colorbar()
                plt.show()
        elif data.ndim == 4:
            num_slices = data.shape[1]  # 第二个维度的大小
            for i in range(num_slices):
                plt.figure()
                plt.title(np.mean(self.p[i,:,:]))
                plt.imshow(data[0, i, :, :], cmap='viridis')  # 绘制第i个切片
                plt.colorbar()
                plt.show()
        else:
            raise ValueError("输入数据的维度错误！")
        
        
        
    def merge(self):
        ds = xr.Dataset()
        var_names=['gph','T','rh','u','v','psfc','slp','u10','v10','T2','tsk','rain','haiwu','tf','gkdbzs','jbzs','sst','dmdz'] #没加dqbd_type
        for var_name in var_names:
            var_data = getattr(self, var_name)  # 从类中获取变量数据
            var_data.attrs['projection'] = ''
            ds[var_name] = var_data  # 添加到 Dataset
        return ds



class swan_ocean:
    def __init__(self,dir_coawst):
        
        roms = xr.open_dataset(os.path.join(dir_coawst,"mangkhut_ocean_his.nc") , engine="netcdf4")
        self.roms = roms
        self.dx,self.dy = self.cal_dxdy()  #m
        # 14海况 共9级；按浪高定义 source:https://baike.baidu.com/item/%E6%B5%B7%E5%86%B5%E8%A1%A8/8027650
        self.hk = xr.DataArray(self.cal_hk(),dims=roms.zeta.dims,coords=roms.zeta.coords)
        # 15海浪
        self.Lwave = roms.Lwave
        self.Lwavep = roms.Lwavep
        self.Dwavep = roms.Dwavep
        self.Pwave_top = roms.Pwave_top
        self.Pwave_bot = roms.Pwave_bot
        self.Uwave_rms = roms.Uwave_rms
        # 19水位
        self.zetaw = roms.zetaw
        # 29海平面高度
        self.zeta = roms.zeta
        # 21海面风  
        self.ssu = roms.Uwind
        self.ssv = roms.Vwind
        # 计算不同标准深度压力，是后面海温，密度的前置量。
        # 这里的p是海水压力，是压力减去一个大气压后的结果; 单位 dbar
        self.z = self.cal_z() #m
        self.p = gsw.p_from_z(self.z,roms.lat_rho.values) 
        # 23海洋标准等深层盐度
        self.salt = roms.salt
        # 22海洋标准等深层海温 
        pt = roms.temp  #等温层位温
        ct = gsw.CT_from_pt(self.salt, pt)   #保守温度
        self.stdT = gsw.t_from_CT(self.salt, ct, self.p)                                                                                           
        # 24海洋标准等深层密度
        self.rho = gsw.rho(self.salt, ct, self.p) 
        # 25海洋标准等深层海流
        self.seau = roms.u
        self.seav = roms.v
        # 26有效波高  #significant height
        self.sh = roms.Hwave
        # 27浪向
        self.lx = roms.Dwave
        # 28海表地转流 sea surface geostrophy wind   source:https://blog.csdn.net/weixin_44237337/article/details/120773341
        ug,vg = self.cal_ssgeowind()
        self.ug = xr.DataArray(ug,dims=roms.zeta.dims,coords=roms.zeta.coords)
        self.vg = xr.DataArray(vg,dims=roms.zeta.dims,coords=roms.zeta.coords)
        # 30海气边界层典型过程
        self.evaporation = roms.evaporation
        self.ssflux = roms.ssflux
        self.shflux = roms.shflux
        # 32海洋内波  vasala frequency
        self.visala = self.cal_visala()
        # 33中尺度涡  OW指数  source:http://www.adearth.ac.cn/article/2020/1001-8166/1001-8166-2020-35-6-568.shtml
        self.ow = xr.DataArray(self.cal_ow(),dims=roms.zeta.dims,coords=roms.zeta.coords)
        # 34海洋锋  用海温梯度表示; 最第15层温度代表海温梯度 ; 0.6°C/100km即出现标准
        #source: 西北太平洋副热带海洋锋的变化特征及其影响
        sstgrd_0m,sstgrd_100m = self.cal_hyf()
        self.sstgrd_0m=xr.DataArray(sstgrd_0m,dims=roms.zeta.dims,coords=roms.zeta.coords)
        self.sstgrd_100m=xr.DataArray(sstgrd_100m,dims=roms.zeta.dims,coords=roms.zeta.coords)
        # 35 跃层强度
        self.wyc_intensity = xr.DataArray( self.cal_wyc(),dims=roms.zeta.dims,coords=roms.zeta.coords )  

        
    def cal_rhogrd(self):
        dz = np.gradient(self.z, axis=1 )
        drho = np.gradient(self.rho ,axis=1)
        return drho/dz
    
    def cal_wyc(self):
        nt,nz,ny,nx = self.z.shape
        #先求垂直梯度
        dT = np.gradient( self.stdT,axis=1  )  # C
        dz = np.gradient( self.z,axis=1 )      #  m
        To_dz= dT/dz*1e3    # c/m
        # 
        wyc_intensity = np.zeros( (nt,ny,nx) )
        for t in range(nt):
            for j in range(ny):
                for i in range(nx):
                    wyc_intensity[t,j,i] = np.max( To_dz[t,:,j,i] )
        return wyc_intensity
        
        
    def cal_hyf(self):
        nt,nz,ny,nx = self.z.shape
        stdT_0m   = np.zeros((nt,ny,nx))
        stdT_100m = np.zeros((nt,ny,nx))
## 并行
        inputs = [(t,self.z[t],self.stdT[t]) for t in range(nt)]  #插值到多高改parallel程序
        if __name__ == "__main__":
            with Pool(24) as pool:
                results = pool.map(parallel,inputs)
            # 把并行处理完的results 变回到原来形状
        for result in results:
            t = result[0]
            stdT_0m[t]   = result[1]  
            stdT_100m[t] = result[2]
# 串行        
        # for t in range(nt):
        #     print(t)
        #     for j in range(ny):
        #         for i in range(nx):
        #             cs           = interp1d(self.z[t,:,j,i], self.stdT[t,:,j,i],fill_value="extrapolate")
        #             stdT_0m[t,j,i]   = cs(0)
        #             stdT_100m[t,j,i] = cs(-100)
        
        
        Tx_0m = np.gradient(stdT_0m, axis=2)/self.dx
        Ty_0m = np.gradient(stdT_0m, axis=1)/self.dy
        Tx_100m = np.gradient(stdT_100m, axis=2)/self.dx
        Ty_100m = np.gradient(stdT_100m, axis=1)/self.dy
        sstgrd_0m   =  np.sqrt( Tx_0m**2 + Ty_0m**2 )*1e3
        sstgrd_100m =  np.sqrt( Tx_100m**2 + Ty_100m**2 )*1e3
        # 测试发现由于海陆交会，sst的差分算法有天然误差；并且有天然边界误差。对异常值做特殊处理。
        sstgrd_0m   = np.where(sstgrd_0m >0.3,0.3,sstgrd_0m )      
        sstgrd_100m = np.where(sstgrd_100m>0.3,0.3,sstgrd_100m)     
        return sstgrd_0m,sstgrd_100m
        
    def cal_ow(self):
        g = 9.81
        f = self.roms.f.values
        h = self.zeta 
        hx = np.gradient(h, axis=2)/self.dx
        hy = np.gradient(h, axis=1)/self.dy
        hxy = np.gradient(hx, axis=1)/self.dy
        hxx = np.gradient(hx, axis=2)/self.dx
        hyy = np.gradient(hy, axis=1)/self.dy
        return 4*g**2*(hxy**2-hxx*hyy)/(f**2)
        
    def cal_dxdy(self):
        R = 6371*1e3  #m
        lon = self.roms.lon_rho.values
        lat = self.roms.lat_rho.values
        dlon = np.gradient( lon , axis=1 )
        dx = R*np.cos(np.radians(lat))*np.radians(dlon)
        dlat = np.gradient(lat , axis=0 )  
        dy = R*np.radians( dlat )
        return dx,dy
        
        
    def cal_visala(self):
        # 计算浮力频率 N
        g = 9.81  
        rho = self.rho 
        d_rho_dz = np.gradient(rho, axis=1)/np.gradient(self.z, axis=1)  # 密度梯度
        N2 = -(g / rho) * d_rho_dz  # 浮力频率平方
        N = np.sqrt(N2)  # 浮力频率
        return N
    
    
    def cal_ssgeowind(self):
        f   = self.roms.f.values
        zeta = self.zeta
        deta_dx = np.gradient(zeta,axis=2)/self.dx # 经向梯度
        deta_dy = np.gradient(zeta,axis=1)/self.dy # 纬向梯度
        # 计算地转流
        g = 9.81
        ug = -(g / f) * deta_dy  # 纬向地转流 (m/s)
        vg = (g / f) * deta_dx  # 经向地转流 (m/s)
        return ug,vg
    
    # 计算不同层次深度
    def cal_z(self):
        h = self.roms.h.values
        hc = self.roms.hc.item()
        s_rho = self.roms.s_rho.values
        Cs_r = self.roms.Cs_r.values
        zeta = self.zeta.values
        # 开始计算
        z = np.zeros(self.roms.salt.shape)
        for k in range(len(s_rho)):
            # 计算每一层的深度
            z[:,k,:,:] = zeta + (zeta + h[np.newaxis, :, :]) * (
                (hc * s_rho[k] + h[np.newaxis, :, :] * Cs_r[k]) / (hc + h[np.newaxis, :, :]))
        return z
        
    
    
    # 利用有效浪高计算海况
    def cal_hk(self):
        a = self.roms.Hwave
        hk = np.zeros( a.shape,dtype=int )
        hk[ (a>=0)&(a<0.1) ]    = 1
        hk[ (a>=0.1)&(a<0.5) ]  = 2
        hk[ (a>=0.5)&(a<1.25) ] = 3
        hk[ (a>=1.25)&(a<2.5) ] = 4
        hk[ (a>=2.5)&(a<4.0) ]  = 5
        hk[ (a>=4.0)&(a<6.0) ]  = 6
        hk[ (a>=6.0)&(a<9.0) ]  = 7
        hk[ (a>=9.0)&(a<14.0) ] = 8
        hk[ a>=14.0 ]           = 9
        return hk
        
        
    def see(self,data):
        if data.ndim == 2:
            plt.contourf(data)
            plt.colorbar()
            plt.show()
        elif data.ndim == 3:
            num_slices = data.shape[0]  # 第一个维度的大小
            for i in range(num_slices):
                plt.figure()
                plt.imshow(data[i, :, :], cmap='viridis')  # 绘制第i个切片
                plt.title(f"{i}")
                plt.colorbar()
                plt.show()
        elif data.ndim == 4:
            num_slices = data.shape[1]  # 第二个维度的大小
            for i in range(num_slices):
                plt.figure()
                plt.imshow(data[0, i, :, :], cmap='viridis')  # 绘制第i个切片
                plt.title(f"{i}")
                plt.colorbar()
                plt.show()
        else:
            raise ValueError("输入数据的维度错误！")
            
    def merge(self):
        """
        将类中的部分变量保存为 NetCDF 文件，并为每个变量添加 unit 和 description 属性。
    
        参数:
        output_path (str): 输出 NetCDF 文件的路径。
        """
        # 创建一个新的 xarray Dataset
        ds = xr.Dataset()
    
        # 定义变量及其对应的单位和描述
        variables = {
            'hk' : {'unit': 'class', 'description': '海浪等级'},
            'Lwave': {'unit': 'm', 'description': '波长'},
            'Lwavep': {'unit': 's', 'description': '波周期'},
            'Dwavep': {'unit': 'degrees', 'description': '波向'},
            'Pwave_top': {'unit': 'Pa', 'description': '波压顶部'},
            'Pwave_bot': {'unit': 'Pa', 'description': '波压底部'},
            'Uwave_rms': {'unit': 'm/s', 'description': '波浪引起的 RMS 流速'},
            'zetaw': {'unit': 'm', 'description': '水位'},
            'zeta': {'unit': 'm', 'description': '海平面高度'},
            'ssu': {'unit': 'm/s', 'description': '表面纬向风速'},
            'ssv': {'unit': 'm/s', 'description': '表面经向风速'},
            'salt': {'unit': 'PSU', 'description': '盐度'},
            'stdT': {'unit': '°C', 'description': '标准等深层海温'},
            'rho': {'unit': 'kg/m^3', 'description': '密度'},
            'seau': {'unit': 'm/s', 'description': '纬向流速'},
            'seav': {'unit': 'm/s', 'description': '经向流速'},
            'sh': {'unit': 'm', 'description': '有效波高'},
            'lx': {'unit': 'degrees', 'description': '浪向'},
            'ug': {'unit': 'm/s', 'description': '纬向地转流'},
            'vg': {'unit': 'm/s', 'description': '经向地转流'},
            'evaporation': {'unit': 'kg/m^2/s', 'description': '蒸发率'},
            'ssflux': {'unit': 'W/m^2', 'description': '表面热通量'},
            'shflux': {'unit': 'W/m^2', 'description': '表面潜热通量'},
            'visala': {'unit': 's^-1', 'description': '浮力频率（Väisälä 频率）'},
            'ow': {'unit': '1', 'description': 'OW 指数（中尺度涡）'},
            'sstgrd_0m':{'unit': '°C/km', 'description': '0m水位密度梯度'},
            'sstgrd_100m':{'unit': '°C/km', 'description': '100m水位密度梯度'},
            'wyc_intensity':{'unit': '°C/km', 'description': '最大垂直温度梯度的水平分布'}
        }
    
        # 将变量添加到 Dataset 中，并设置属性
        for var_name, var_attrs in variables.items():
            var_data = getattr(self, var_name)  # 从类中获取变量数据
            ds[var_name] = var_data  # 添加到 Dataset
            ds[var_name].attrs['unit'] = var_attrs['unit']  # 添加单位属性
            ds[var_name].attrs['description'] = var_attrs['description']  # 添加描述属性
        return ds

        
postwrf = wrfout(wrfoutpath).merge()
postos = swan_ocean(dir_coawst).merge()
post = xr.merge([postwrf,postos])
post.to_netcdf('./postcoawst.nc')


def see(data):
    if data.ndim == 2:
        plt.contourf(data)
        plt.colorbar()
        plt.show()
    elif data.ndim == 3:
        num_slices = data.shape[0]  # 第一个维度的大小
        for i in range(num_slices):
            plt.figure()
            plt.imshow(data[i, :, :], cmap='viridis')  # 绘制第i个切片
            plt.colorbar()
            plt.show()
    elif data.ndim == 4:
        num_slices = data.shape[1]  # 第二个维度的大小
        for i in range(num_slices):
            plt.figure()
            plt.imshow(data[0, i, :, :], cmap='viridis')  # 绘制第i个切片
            plt.colorbar()
            plt.show()
    else:
        raise ValueError("输入数据的维度错误！")

